#include "yolo_threadpool.h"

YoloThreadPool::YoloThreadPool(const std::string model_path, const int thread_num,
                               const std::string label_path)
    : model_path_(model_path), thread_num_(thread_num), label_path_(label_path)
{
    Init();
}

YoloThreadPool::~YoloThreadPool()
{
    DeInint();
}

void YoloThreadPool::Init()
{
    init_post_process(label_path_);

    try
    {
        // 配置线程池
        pool_ = std::make_unique<ThreadPool>(thread_num_);

        // 每一个线程加载一个模型
        for (int i = 0; i < thread_num_; i++)
        {
            models_.push_back(std::make_shared<Yolov8Custom>(std::forward<std::string>(model_path_)));
        }
    }
    catch (const std::bad_alloc &e)
    {
        NN_LOG_ERROR("ThreadPool init failed: %s", e.what());
        exit(EXIT_FAILURE);
    }

    for (int i = 0; i < thread_num_; i++)
    {
        auto ret = models_[i]->Init(models_[0]->get_rknn_context(), i != 0);

        if (ret != RKNN_SUCC)
        {
            NN_LOG_ERROR("Model init failed");
            exit(EXIT_FAILURE);
        }
    }
}

void YoloThreadPool::DeInint() { deinit_post_process(); }

void YoloThreadPool::AddInferenceTask(std::shared_ptr<cv::Mat> src, Preprocess &preprocess)
{
    pool_->enqueue(
        [&](std::shared_ptr<cv::Mat> original_img)
        {
            auto convert_img = preprocess.Convert(*original_img);
            auto mode_id = get_model_id();
            cv::Mat rgb_img = cv::Mat::zeros(
                this->models_[mode_id]->get_model_width(),
                this->models_[mode_id]->get_model_height(), convert_img->type());

            cv::cvtColor(*convert_img, rgb_img, cv::COLOR_BGR2RGB);
            object_detect_result_list od_results;

            models_[mode_id]->Inference(rgb_img.ptr(), &od_results, preprocess.get_letter_box());

            preprocess.ImagePostProcess(*original_img, od_results);

            std::lock_guard<std::mutex> lock_guard(image_results_mutex_);
            image_results_.push(std::move(original_img));
        },
        std::move(src));
}

int YoloThreadPool::get_model_id()
{
    std::lock_guard<std::mutex> lock_guard(id_mutex_);
    int model_id = id;

    id++;

    if (id == thread_num_)
    {
        id = 0;
    }

    return model_id;
}

std::shared_ptr<cv::Mat> YoloThreadPool::GetImageResultFromQueue()
{
    std::lock_guard<std::mutex> lock_guard(image_results_mutex_);
    if (this->image_results_.empty())
    {
        return nullptr;
    }

    else
    {
        auto res = this->image_results_.front();
        this->image_results_.pop();
        return std::move(res);
    }
}

int YoloThreadPool::GetTaseksSize()
{
    // std::lock_guard<std::mutex> lock_guard(image_results_mutex_);
    return pool_->TasksSize();
}